package com.example.scheduler;

import java.util.List;
import java.util.concurrent.ConcurrentHashMap;

/**
 * DagTask：表示一次DAG任务执行实例。
 * levels中保存若干层；taskStatus保存“dagId:子任务” -> 是否完成。
 */
public class DagTask {
    private final String dagId;
    private final List<List<String>> levels;
    private final ConcurrentHashMap<String, Boolean> taskStatus;
    private volatile int currentLevel;

    public DagTask(String dagId, List<List<String>> levels) {
        this.dagId = dagId;
        this.levels = levels;
        this.currentLevel = 0;
        this.taskStatus = new ConcurrentHashMap<>();
        // 初始化第一层任务的执行状态
        if (!levels.isEmpty()) {
            for (String taskId : levels.get(0)) {
                taskStatus.put(dagId + ":" + taskId, false);
            }
        }
    }

    public String getDagId() {
        return dagId;
    }

    public List<String> getCurrentLevelTasks() {
        if (currentLevel < levels.size()) {
            return levels.get(currentLevel);
        }
        return null;
    }

    public void markTaskFinished(String fullTaskId) {
        taskStatus.put(fullTaskId, true);
    }

    public boolean isCurrentLevelComplete() {
        List<String> currentTasks = levels.get(currentLevel);
        for (String t : currentTasks) {
            String fullId = dagId + ":" + t;
            if (!taskStatus.getOrDefault(fullId, false)) {
                return false;
            }
        }
        return true;
    }

    public boolean moveToNextLevel() {
        if (currentLevel + 1 < levels.size()) {
            currentLevel++;
            for (String taskId : levels.get(currentLevel)) {
                taskStatus.put(dagId + ":" + taskId, false);
            }
            return true;
        }
        return false;
    }

    public boolean isCompleted() {
        // 只有当 currentLevel >= levels.size() 并且最后一层完成，才算DAG结束
        return (currentLevel >= levels.size()) && isCurrentLevelComplete();
    }
}